from langchain_ollama.llms import OllamaLLM
import yaml
from langchain.agents.format_scratchpad import format_to_openai_function_messages
from langchain_community.utilities import SerpAPIWrapper
from langchain_community.agent_toolkits.load_tools import load_tools
from langchain.agents import create_react_agent,initialize_agent,create_json_agent
from langchain.agents import AgentType
from langchain.agents.mrkl.output_parser import MRKLOutputParser
from langchain.prompts import PromptTemplate
from langchain_core.globals import set_debug
import logging

# 设置日志记录级别为 WARNING 或 ERROR
logging.basicConfig(level=logging.WARNING)
with open('config.yaml', 'r') as file:
    data = yaml.safe_load(file)


modelOllama = OllamaLLM(model="llama3:latest")


set_debug(True)
print(data['serpapi_api_key'])


# from langchain.output_parsers import StructuredOutputParser, ResponseSchema

# # 定义我们想要接收的响应模式
# response_schemas = [
#     ResponseSchema(name="question", description="提问的问题"),
#     ResponseSchema(name="description", description="问题的描述文案"),
#     ResponseSchema(name="answer", description="问题的回答答案"),
# ]
# # 创建输出解析器
# output_parser = StructuredOutputParser.from_response_schemas(response_schemas)


# class Tool:
#     def __init__(self, name: str, description: str):
#         self.name = name
#         self.description = description

# # 创建一些工具实例
# tool1 = Tool("SearchEngine", "A tool for searching the web.")
# tool2 = Tool("DatabaseLookup", "A tool for looking up data in a database.")

# tools = [tool1,tool2]

template = '''Answer the following questions as best you can. 

You have access to the following tools:{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer includes the return_value, and the return_value is the answer 

Begin!

Question: {input}
Thought:{agent_scratchpad}'''
from langchain.tools import Tool
# from langchain_core.tools import BaseTool
prompt = PromptTemplate.from_template(template)
search = SerpAPIWrapper(serpapi_api_key=data['serpapi_api_key'])
tools = [
    Tool(
        name="Search",
        func=search.run,
        description="useful for when you need to answer questions about current events"
    ),
]
# tools=[]
agent = create_react_agent(
    modelOllama, 
    tools,
    prompt
)
# intermediate_steps = []

# res=agent.invoke({"input": "do you know the leetcode?",
#               "intermediate_steps": intermediate_steps,
# })

# print("res=",res)

# 导入AgentExecutor
from langchain.agents import AgentExecutor
# 创建代理执行器并传入代理和工具
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True,handle_parsing_errors=True)
# 调用代理执行器，传入输入数据
print("第一次运行的结果：")
agent_executor.invoke({"input": "你知道今天是几月几号吗，历史上的今天发生了什么?"})
# print("第二次运行的结果：")
# agent_executor.invoke({"input": "你知道联影智元和联影医疗的关系吗?"})
